# -*- coding: utf-8 -*-

"""
基于 Pytdx API 的数据任务基类

借鉴alphahome的TushareTask实现，为所有Pytdx任务提供通用功能。
"""

import abc
import asyncio
import logging
import os
import math
import time
from datetime import date
from typing import Any, Dict, List, Optional, Set

import pandas as pd

from alphahome.fetchers.base.fetcher_task import FetcherTask
from .pytdx_api import PytdxAPI
from .pytdx_data_transformer import PytdxDataTransformer


class PytdxConnectionPool:
    """
    Pytdx连接池管理器

    借鉴hikyuu的连接复用思想，实现连接池以提高性能和稳定性。
    提供连接的自动管理、健康检查和资源控制。
    """

    def __init__(self, max_connections: int = 3, host: Optional[str] = None,
                 port: int = 7709, timeout: int = 10, health_check_interval: int = 30):
        """
        初始化连接池

        Args:
            max_connections: 最大连接数，默认3个（pytdx并发限制）
            host: 通达信服务器地址
            port: 端口号
            timeout: 连接超时时间
            health_check_interval: 健康检查间隔（秒）
        """
        self.max_connections = max_connections
        self.host = host
        self.port = port
        self.timeout = timeout
        self.health_check_interval = health_check_interval

        # 连接池队列和状态管理
        self._pool: asyncio.Queue[PytdxAPI] = asyncio.Queue()
        self._active_connections: Set[PytdxAPI] = set()
        self._connection_stats: Dict[PytdxAPI, Dict] = {}

        # 健康监控任务
        self._health_monitor_task: Optional[asyncio.Task] = None
        self._is_monitoring = False

        # 统计信息
        self._total_connections_created = 0
        self._total_connections_reused = 0
        self._total_connections_failed = 0

    async def start_health_monitoring(self):
        """启动健康监控"""
        if not self._is_monitoring:
            self._is_monitoring = True
            self._health_monitor_task = asyncio.create_task(self._health_monitor_loop())
            self.logger.info(f"Pytdx连接池健康监控已启动，检查间隔: {self.health_check_interval}秒")

    async def stop_health_monitoring(self):
        """停止健康监控"""
        if self._is_monitoring:
            self._is_monitoring = False
            if self._health_monitor_task:
                self._health_monitor_task.cancel()
                try:
                    await self._health_monitor_task
                except asyncio.CancelledError:
                    pass
            self.logger.info("Pytdx连接池健康监控已停止")

    async def get_connection(self) -> PytdxAPI:
        """
        获取可用连接

        Returns:
            可用的PytdxAPI连接实例
        """
        try:
            # 首先尝试从池中获取现有连接
            api = self._pool.get_nowait()
            if await self._is_connection_alive(api):
                self._total_connections_reused += 1
                self.logger.debug(f"复用现有连接，总复用次数: {self._total_connections_reused}")
                return api
            else:
                # 连接失效，清理并创建新连接
                await self._cleanup_connection(api)
                self.logger.debug("清理失效连接，准备创建新连接")

        except asyncio.QueueEmpty:
            pass

        # 创建新连接或等待可用连接
        if len(self._active_connections) < self.max_connections:
            return await self._create_new_connection()
        else:
            # 等待池中有可用连接
            self.logger.debug("连接池已满，等待可用连接...")
            api = await self._pool.get()
            return api

    async def return_connection(self, api: PytdxAPI):
        """
        归还连接到池中

        Args:
            api: 要归还的连接
        """
        if api in self._active_connections:
            # 更新连接统计信息
            if api in self._connection_stats:
                self._connection_stats[api]['last_used'] = time.time()
                self._connection_stats[api]['use_count'] += 1

            # 重新放回池中
            await self._pool.put(api)
            self.logger.debug(f"连接已归还到池中，当前池大小: {self._pool.qsize()}")

    async def _create_new_connection(self) -> PytdxAPI:
        """创建新连接"""
        try:
            api = PytdxAPI(host=self.host, port=self.port, timeout=self.timeout)
            success = await api.connect()

            if success:
                self._active_connections.add(api)
                self._connection_stats[api] = {
                    'created_at': time.time(),
                    'last_used': time.time(),
                    'use_count': 0,
                    'error_count': 0
                }
                self._total_connections_created += 1

                self.logger.info(f"成功创建新连接，总连接数: {len(self._active_connections)}/{self.max_connections}")
                return api
            else:
                self._total_connections_failed += 1
                raise Exception("连接通达信服务器失败")

        except Exception as e:
            self.logger.error(f"创建新连接失败: {e}")
            raise

    async def _is_connection_alive(self, api: PytdxAPI) -> bool:
        """
        检查连接是否仍然有效

        Args:
            api: 要检查的连接

        Returns:
            连接是否有效
        """
        try:
            # 发送轻量级测试请求
            test_data = await api.get_stock_daily_bars('SH', '000001', start=0, count=1)
            return test_data is not None and len(test_data) > 0
        except Exception as e:
            self.logger.debug(f"连接健康检查失败: {e}")
            return False

    async def _cleanup_connection(self, api: PytdxAPI):
        """清理失效连接"""
        if api in self._active_connections:
            self._active_connections.remove(api)
            self._connection_stats.pop(api, None)

        try:
            api.disconnect()
        except:
            pass

        self.logger.debug("已清理失效连接")

    async def _health_monitor_loop(self):
        """健康监控循环"""
        while self._is_monitoring:
            try:
                await asyncio.sleep(self.health_check_interval)

                unhealthy_connections = []
                for api in list(self._active_connections):
                    if not await self._is_connection_alive(api):
                        unhealthy_connections.append(api)

                # 清理不健康的连接
                for api in unhealthy_connections:
                    self.logger.warning("检测到不健康连接，正在清理并重建")
                    await self._cleanup_connection(api)

                    # 尝试创建新连接来替换
                    try:
                        await self._create_new_connection()
                    except Exception as e:
                        self.logger.error(f"重建连接失败: {e}")

            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"健康监控循环异常: {e}")

    async def close_all_connections(self):
        """关闭所有连接并清理资源"""
        await self.stop_health_monitoring()

        # 关闭所有活跃连接
        for api in list(self._active_connections):
            try:
                api.disconnect()
            except:
                pass

        self._active_connections.clear()
        self._connection_stats.clear()

        # 清空连接池
        while not self._pool.empty():
            try:
                self._pool.get_nowait()
            except asyncio.QueueEmpty:
                break

        self.logger.info("所有连接已关闭，连接池已清理")

    def get_stats(self) -> Dict[str, Any]:
        """获取连接池统计信息"""
        return {
            'active_connections': len(self._active_connections),
            'pool_size': self._pool.qsize(),
            'max_connections': self.max_connections,
            'total_created': self._total_connections_created,
            'total_reused': self._total_connections_reused,
            'total_failed': self._total_connections_failed,
            'connection_stats': {
                f"conn_{i}": stats
                for i, (conn, stats) in enumerate(self._connection_stats.items())
            }
        }

    @property
    def logger(self):
        """获取日志器"""
        return logging.getLogger(f"{__name__}.ConnectionPool")


class PytdxTask(FetcherTask, abc.ABC):
    """
    一个针对 Pytdx API 的抽象任务基类。

    此类继承自 FetcherTask，并实现了 Pytdx 特有的技术逻辑：
    - 在 `fetch_batch` 中统一调用 Pytdx API 并进行标准数据转换。
    - 提供一个通用的 `prepare_params` 实现。

    它将 `get_batch_list` 继续声明为抽象方法，因为每个具体的 Pytdx 接口
    都有其独特的批处理要求。
    """

    data_source = "pytdx"

    # Pytdx 特有配置
    default_host: Optional[str] = None
    default_port: int = 7709
    default_timeout: int = 10

    def __init__(self, db_connection, host=None, port=None, api=None,
                 connection_pool=None, use_connection_pool=True, **kwargs):
        """
        初始化 PytdxTask。

        Args:
            db_connection: 数据库连接。
            host (str, optional): Pytdx服务器地址。
            port (int, optional): Pytdx服务器端口。
            api (PytdxAPI, optional): 已初始化的 PytdxAPI 实例（传统模式）。
            connection_pool (PytdxConnectionPool, optional): 连接池实例。
            use_connection_pool (bool, optional): 是否使用连接池模式，默认True。
            **kwargs: 传递给 FetcherTask 的参数。
        """
        # 从 kwargs 中提取 task_config，以便 _apply_config 能正确工作
        task_config = kwargs.get("task_config", {})

        # 将 Pytdx 特有的配置添加到 task_config 中
        if "host" not in task_config and host is None:
            task_config["host"] = self.default_host
        if "port" not in task_config and port is None:
            task_config["port"] = self.default_port
        if "timeout" not in task_config:
            task_config["timeout"] = self.default_timeout

        kwargs["task_config"] = task_config

        super().__init__(db_connection, **kwargs)

        # 初始化连接管理
        self.host = host or task_config.get("host")
        self.port = port or task_config.get("port", self.default_port)
        self.timeout = task_config.get("timeout", self.default_timeout)

        # 连接模式决策逻辑
        if connection_pool is not None:
            # 显式提供了连接池，使用连接池模式
            self.connection_pool = connection_pool
            self.api = None
            self._use_connection_pool = True
            self._is_default_connection_pool = False
        elif api is not None:
            # 显式提供了API，使用传统模式
            self.api = api
            self.connection_pool = None
            self._use_connection_pool = False
            self._is_default_connection_pool = False
        elif use_connection_pool:
            # 默认使用连接池模式，自动创建连接池
            self.connection_pool = PytdxConnectionPool(
                max_connections=3,
                host=self.host,
                port=self.port,
                timeout=self.timeout,
                health_check_interval=30
            )
            self.api = None
            self._use_connection_pool = True
            self._is_default_connection_pool = True
            self.logger.info("PytdxTask: 使用默认连接池模式（3个连接）")
        else:
            # 传统模式：创建单个连接
            self.api = PytdxAPI(
            host=self.host,
            port=self.port,
            timeout=self.timeout
        )
            self.connection_pool = None
            self._use_connection_pool = False
            self._is_default_connection_pool = False

        self.data_transformer = PytdxDataTransformer(self)

    def _calculate_smart_batch_size(self, last_update_date: str = None, data_type: str = "daily") -> int:
        """
        智能计算批次大小，借鉴hikyuu的guess_day_n_step实现

        Args:
            last_update_date: 最后更新日期，格式YYYYMMDD
            data_type: 数据类型 ('daily', '1min', '5min', 'weekly', 'monthly')

        Returns:
            推荐的批次大小
        """
        try:
            today = date.today()

            if not last_update_date:
                # 没有历史数据，使用最大批次
                return 800

            # 解析最后更新日期
            last_date = date(
                int(last_update_date[:4]),
                int(last_update_date[4:6]),
                int(last_update_date[6:8])
            )

            # 计算相差的天数
            days_diff = (today - last_date).days

            if days_diff <= 0:
                # 数据已是最新，只获取少量数据验证
                return 10
            elif days_diff <= 30:
                # 最近一个月的数据，使用较小的批次
                return min(days_diff + 5, 800)
            elif days_diff <= 365:
                # 最近一年的数据，使用中等批次
                return min(int(days_diff * 0.8), 800)
            else:
                # 历史较久的数据，使用最大批次
                # 估算年份差，考虑交易日大约250天
                years_diff = (today.year - last_date.year + 1)

                # 根据数据类型调整估算逻辑
                if data_type == "daily":
                    estimated_trading_days = years_diff * 250
                elif data_type in ["1min", "5min"]:
                    # 分钟数据估算更保守
                    estimated_trading_days = years_diff * 200
                else:
                    estimated_trading_days = years_diff * 250

                estimated_batches = int(estimated_trading_days // 800)
                if estimated_batches <= 1:
                    # 计算精确的天数（考虑可能的交易日）
                    return min(days_diff, 800)
                else:
                    return 800

        except Exception as e:
            self.logger.warning(f"智能批次计算失败，使用默认值: {e}")
            return 800

    def _validate_data_integrity(self, new_data: pd.DataFrame, existing_data: pd.DataFrame = None) -> bool:
        """
        验证数据完整性，借鉴hikyuu的前后数据对比机制

        Args:
            new_data: 新获取的数据
            existing_data: 现有的历史数据

        Returns:
            数据是否完整有效
        """
        if new_data.empty:
            return True  # 空数据视为有效

        try:
            # 检查新数据的基本合理性
            if not self._validate_basic_data_rules(new_data):
                return False

            # 如果有历史数据，进行前后对比校验
            if existing_data is not None and not existing_data.empty:
                return self._validate_continuity_with_existing(new_data, existing_data)

            return True

        except Exception as e:
            self.logger.error(f"数据完整性校验失败: {e}")
            return False

    def _validate_basic_data_rules(self, df: pd.DataFrame) -> bool:
        """
        验证基本数据规则
        """
        try:
            # 价格必须为正数
            price_checks = [
                (df["close"] > 0, "收盘价必须为正数"),
                (df["open"] > 0, "开盘价必须为正数"),
                (df["high"] > 0, "最高价必须为正数"),
                (df["low"] > 0, "最低价必须为正数"),
            ]

            for condition, message in price_checks:
                if not condition.all():
                    invalid_count = (~condition).sum()
                    self.logger.warning(f"数据完整性校验失败: {message}, 发现{invalid_count}条无效记录")
                    return False

            # 成交量和成交额不能为负数
            volume_checks = [
                (df["volume"] >= 0, "成交量不能为负数"),
                (df["amount"] >= 0, "成交额不能为负数"),
            ]

            for condition, message in volume_checks:
                if not condition.all():
                    invalid_count = (~condition).sum()
                    self.logger.warning(f"数据完整性校验失败: {message}, 发现{invalid_count}条无效记录")
                    return False

            # 价格逻辑关系检查
            logic_checks = [
                (df["high"] >= df["low"], "最高价不能低于最低价"),
                (df["high"] >= df["open"], "最高价不能低于开盘价"),
                (df["high"] >= df["close"], "最高价不能低于收盘价"),
                (df["low"] <= df["open"], "最低价不能高于开盘价"),
                (df["low"] <= df["close"], "最低价不能高于收盘价"),
            ]

            for condition, message in logic_checks:
                if not condition.all():
                    invalid_count = (~condition).sum()
                    self.logger.warning(f"数据完整性校验失败: {message}, 发现{invalid_count}条无效记录")
                    return False

            return True

        except Exception as e:
            self.logger.error(f"基本数据规则校验失败: {e}")
            return False

    def _validate_continuity_with_existing(self, new_data: pd.DataFrame, existing_data: pd.DataFrame) -> bool:
        """
        验证新数据与现有数据的连续性
        """
        try:
            # 确保数据按日期排序
            new_data_sorted = new_data.sort_values('trade_date')
            existing_data_sorted = existing_data.sort_values('trade_date')

            # 获取现有数据的最后一条记录
            last_existing = existing_data_sorted.iloc[-1]

            # 在新数据中查找相同日期的记录
            same_date_mask = new_data_sorted['trade_date'] == last_existing['trade_date']
            if same_date_mask.any():
                # 找到相同日期的记录，进行精确比对
                matching_record = new_data_sorted[same_date_mask].iloc[0]

                # 比对关键字段（允许小幅误差）
                tolerance = 0.02  # 2%的容差

                price_fields = ['open', 'high', 'low', 'close']
                for field in price_fields:
                    existing_val = last_existing[field]
                    new_val = matching_record[field]

                    if existing_val > 0:  # 避免除零错误
                        diff_ratio = abs(existing_val - new_val) / existing_val
                        if diff_ratio > tolerance:
                            self.logger.error(
                                f"数据连续性校验失败: {field}字段不匹配 "
                                f"现有数据: {existing_val}, 新数据: {new_val}, "
                                f"差异比例: {diff_ratio:.2%}"
                            )
                            return False

                # 比对成交量和成交额（相对宽松的检查）
                volume_tolerance = 0.1  # 10%的容差
                volume_fields = ['volume', 'amount']
                for field in volume_fields:
                    existing_val = last_existing[field]
                    new_val = matching_record[field]

                    if existing_val > 0:
                        diff_ratio = abs(existing_val - new_val) / existing_val
                        if diff_ratio > volume_tolerance:
                            self.logger.warning(
                                f"数据连续性校验警告: {field}字段差异较大 "
                                f"现有数据: {existing_val}, 新数据: {new_val}, "
                                f"差异比例: {diff_ratio:.2%}"
                            )
                            # 对于成交量/额，警告但不阻断

            return True

        except Exception as e:
            self.logger.error(f"数据连续性校验失败: {e}")
            return False

    async def _get_last_update_date(self, market: str, code: str) -> str:
        """
        获取股票的最后更新日期

        Args:
            market: 市场代码
            code: 股票代码

        Returns:
            最后更新日期，格式YYYYMMDD，如果没有数据返回None
        """
        try:
            ts_code = f"{code}.{market.upper()}"

            # 从数据库查询最后更新日期
            query = f"""
                SELECT MAX({self.date_column}) as last_date
                FROM {self.table_name}
                WHERE ts_code = $1
            """

            result = await self.db_connection.fetchval(query, ts_code)

            if result:
                # 确保返回字符串格式
                if isinstance(result, str):
                    return result
                elif hasattr(result, 'strftime'):
                    return result.strftime('%Y%m%d')
                else:
                    return str(result).replace('-', '')

            return None

        except Exception as e:
            self.logger.debug(f"获取 {market}{code} 最后更新日期失败: {e}")
            return None

    def _get_existing_data_for_validation(self, market: str, code: str) -> pd.DataFrame:
        """
        获取现有数据用于完整性校验

        Args:
            market: 市场代码
            code: 股票代码

        Returns:
            现有的历史数据DataFrame
        """
        try:
            ts_code = f"{code}.{market.upper()}"

            # 从数据库获取最近的一些历史数据用于校验
            query = f"""
                SELECT ts_code, trade_date, open, high, low, close, volume, amount
                FROM {self.table_name}
                WHERE ts_code = $1
                ORDER BY trade_date DESC
                LIMIT 10
            """

            # 注意：这里需要同步查询，因为在_transform_data中无法使用async
            # 在实际应用中，可能需要调整架构或者缓存数据
            import asyncio
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

            try:
                rows = loop.run_until_complete(self.db_connection.fetch(query, ts_code))
                if rows:
                    return pd.DataFrame([dict(row) for row in rows])
            finally:
                loop.close()

            return pd.DataFrame()

        except Exception as e:
            self.logger.debug(f"获取 {market}{code} 现有数据失败: {e}")
            return pd.DataFrame()

    def _apply_config(self, task_config: Dict):
        """合并代码默认值和配置文件设置。"""
        super()._apply_config(task_config)  # 调用父类的配置应用

        cls = type(self)
        self.host = task_config.get("host", cls.default_host)
        self.port = int(task_config.get("port", cls.default_port))
        self.timeout = int(task_config.get("timeout", cls.default_timeout))

    async def ensure_connection(self) -> bool:
        """确保与Pytdx服务器的连接（向后兼容）"""
        if not self._use_connection_pool and self.api:
            # 传统模式：检查单个连接
            return self.api.connected
        return True  # 连接池模式总是返回True

    async def prepare_params(self, batch_params: Dict) -> Dict:
        """
        准备 Pytdx API 请求的参数。

        Args:
            batch_params: 批次参数字典。

        Returns:
            准备好的参数字典。
        """
        # Pytdx的参数相对简单，主要处理市场和代码信息
        prepared_params = {}

        # 复制所有批次参数
        prepared_params.update(batch_params)

        # 确保必要的参数存在
        if 'market' not in prepared_params:
            prepared_params['market'] = 'SH'  # 默认上海市场

        if 'code' not in prepared_params:
            raise ValueError("批次参数中必须包含 'code'")

        return prepared_params

    async def fetch_batch(self, batch_params: Dict, stop_event: Optional[asyncio.Event] = None) -> pd.DataFrame:
        """
        从 Pytdx API 获取一批数据。

        支持连接池和传统单连接两种模式。

        Args:
            batch_params: 批次参数。
            stop_event: 取消事件。

        Returns:
            获取到的数据 DataFrame。
        """
        api = None
        try:
            # 检查是否被取消
            if stop_event and stop_event.is_set():
                self.logger.warning("任务被取消")
                raise asyncio.CancelledError("任务在fetch_batch中被取消")

            # 获取连接（连接池模式或传统模式）
            if self._use_connection_pool:
                api = await self.connection_pool.get_connection()
            else:
                # 传统模式：确保连接
                if not await self.ensure_connection():
                    self.logger.error("无法连接到Pytdx服务器")
                    return pd.DataFrame()
                    api = self.api

            # 检查是否被取消
            if stop_event and stop_event.is_set():
                self.logger.warning("任务被取消")
                raise asyncio.CancelledError("任务在连接后被取消")

            # 准备参数
            params = await self.prepare_params(batch_params)

            # 调用具体的获取方法（由子类实现）
            raw_data = await self._fetch_raw_data_with_api(api, params)

            if raw_data is None:
                return pd.DataFrame()

            # 转换数据格式
            df = self._transform_data(raw_data, params)

            self.logger.debug(f"成功获取 {len(df)} 条数据记录")
            return df

        except Exception as e:
            self.logger.error(f"获取批次数据时出错: {e}", exc_info=True)
            return pd.DataFrame()
        finally:
            # 连接池模式：归还连接
            if self._use_connection_pool and api is not None:
                await self.connection_pool.return_connection(api)

    async def _fetch_raw_data_with_api(self, api: PytdxAPI, params: Dict) -> Optional[List[Dict]]:
        """
        使用指定的API连接获取原始数据

        这个方法提供了统一的接口，让子类可以通过传入的API实例进行数据获取。
        支持连接池模式和传统模式。

        Args:
            api: PytdxAPI连接实例
            params: 请求参数

        Returns:
            原始数据列表
        """
        # 默认实现：调用子类的_fetch_raw_data方法
        # 子类可以重写这个方法来直接使用api参数
        return await self._fetch_raw_data(params)

    @abc.abstractmethod
    async def _fetch_raw_data(self, params: Dict) -> Optional[List[Dict]]:
        """
        从Pytdx API获取原始数据（由子类实现）

        Args:
            params: 准备好的参数

        Returns:
            原始数据列表
        """
        pass

    @abc.abstractmethod
    def _transform_data(self, raw_data: List[Dict], params: Dict) -> pd.DataFrame:
        """
        转换原始数据为标准格式（由子类实现）

        Args:
            raw_data: 原始数据
            params: 参数信息

        Returns:
            标准格式的DataFrame
        """
        pass

    async def cleanup(self):
        """清理资源"""
        if self._use_connection_pool:
            # 连接池模式：停止健康监控
            if self.connection_pool:
                await self.connection_pool.stop_health_monitoring()
                # 如果是默认创建的连接池，完全关闭它
                if getattr(self, '_is_default_connection_pool', False):
                    await self.connection_pool.close_all_connections()
                    self.logger.info("PytdxTask: 默认连接池已完全关闭")
        else:
            # 传统模式：断开单个连接
            if self.api:
                self.api.disconnect()
            self._connected = False
